import qualified Data.Map as Map 
import Data.List (isPrefixOf)
import Data.Char (isDigit)

data MaskKind = Reset | Set | Keep deriving (Show)
type Mask = [MaskKind]

type Address = Integer
type Value = Integer

data Instruction = SetMask Mask 
                 | WriteMem Address Value deriving (Show)

type Memory = Map.Map Address Value
data Computer = Computer Mask Memory deriving (Show)

parseMaskKind :: Char -> MaskKind
parseMaskKind c = case c of
    '0' -> Reset
    '1' -> Set
    'X' -> Keep

-- mask = X11001110001101XX01111X1001X01101111
-- mem[32163] = 23587
parseInstruction :: String -> Instruction
parseInstruction string 
    | isPrefixOf "mask" string = SetMask mask
    | isPrefixOf "mem" string  = WriteMem address value
    | otherwise = error "Unkonw instruction!"
        where
            mask = map parseMaskKind . drop 7 $ string
            address = read . takeWhile isDigit . drop 4 $ string
            value = read . drop 2 . dropWhile (/='=') $ string

parseInput :: String -> [Instruction]
parseInput = map parseInstruction . lines

binaryToValue :: String -> Value
binaryToValue = foldl foldingFunc 0
    where
        foldingFunc :: Value -> Char -> Value
        foldingFunc value ch = case ch of
                '0' -> value * 2
                '1' -> value * 2 + 1

valueToBinary :: Value -> String
valueToBinary value = (replicate len '0') ++ string
    where
        string = reverse . valueToBinaryRev $ value
        len = 36 - length string
    
        valueToBinaryRev :: Value -> String
        valueToBinaryRev 1 = "1"
        valueToBinaryRev 0 = "0"
        valueToBinaryRev value = (if reminder == 1 then '1' else '0') : valueToBinaryRev half
            where
                reminder = value `mod` 2
                half = value `div` 2

withMask :: MaskKind -> Char -> Char
withMask m c = case m of
    Reset -> '0'
    Set -> '1'
    Keep -> c

maskValue :: Mask -> Value -> Value
maskValue mask value = binaryToValue binStr'
    where
        binStr = valueToBinary value
        binStr' = zipWith withMask mask binStr

executeInstruction :: Computer -> Instruction -> Computer
executeInstruction (Computer _ memory) (SetMask mask) = Computer mask memory
executeInstruction (Computer mask memory) (WriteMem address value) = Computer mask memory'
    where
        value' = maskValue mask value
        memory' = Map.insert address value' memory

execute :: [Instruction] -> Computer 
execute = foldl executeInstruction (Computer (repeat Keep) Map.empty)

memorySum :: Computer -> Value
memorySum (Computer _ memory) = sum . map snd . Map.toList $ memory

main :: IO ()
main = do
    input <- readFile "input.txt"
    print . memorySum . execute . parseInput $ input